import torch.nn as nn

from layers import HGCN, AvgReadout, Discriminator


class HDGI(nn.Module):
    def __init__(self, n_input, hidden, semantic_hidden, p, activation):
        super(HDGI, self).__init__()
        self.hgcn = HGCN(n_input, hidden, semantic_hidden, p, activation)
        self.read = AvgReadout()
        self.sigmoid = nn.Sigmoid()
        self.disc = Discriminator(hidden)

    def forward(self, pos, neg, adj, sparse, msk, samp_bias1, samp_bias2):
        h_1 = self.hgcn(pos, adj, sparse)

        c = self.read(h_1, msk)
        c = self.sigmoid(c)

        h_2 = self.hgcn(neg, adj, sparse)

        ret = self.disc(c, h_1, h_2, samp_bias1, samp_bias2)
        return ret

    def embed(self, seq, adj, sparse, msk):
        h_1 = self.hgcn(seq, adj, sparse)
        c = self.read(h_1, msk)

        return h_1.detach(), c.detach()
